#!/usr/bin/python

# Protocol1
# Workflow:
# 1. BLASTS input sequence and gets gis
# 2. Runs gis through entrez to annotate
# 3. Puts sequences through make_table5 to filter

# Contact kjwu@ucsd.edu for questions.

# Reference:
# http://www.biopython.org/wiki/Annotate_Entrez_Gene_IDs
# http://www.biopython.org/DIST/docs/tutorial/Tutorial.html#htoc51
# http://www.biopython.org/DIST/docs/tutorial/Tutorial.html#htoc101
# http://www.biopython.org/DIST/docs/tutorial/Tutorial.html#htoc105
# http://eutils.ncbi.nlm.nih.gov/corehtml/query/static/efetchseq_help.html

from Bio import Entrez
from Bio import SeqIO
from Bio.Blast import NCBIWWW
from Bio.Blast import NCBIXML
import urllib2
import subprocess
import tempfile

# success:              whether function completed normally
# blast_in:             input sequence/number
# blast_threshold:      threshold for blast e value
# make_table_threshold: threshold for make_table5
# gi_num:               blast'd gi numbers
# make_table_out:       stdout and stderr from make_table5 run
# faa_file:             file handle to make_table5 result
class protocol1:
  Entrez.email = 'kjwu@ucsd.edu'

  # params: input to be blasted, blast_threshold, make_table_threshold
  def __init__(self, blast_in, blast_threshold=0.005, make_table_threshold=0.7):

    self.success = False
    self.errors = None

    # saves inputs for future reference
    self.blast_in = blast_in
    self.blast_threshold = blast_threshold
    self.make_table_threshold = make_table_threshold

    self.gi_num = []

    # run ncbi blast, parse results and write gis to gi_num file
    try:
      result_handle = NCBIWWW.qblast("blastp", "nr", blast_in, format_type="XML", expect=blast_threshold)
      
      for blast_record in NCBIXML.parse(result_handle):
        for alignment in blast_record.alignments:
          self.gi_num.append(alignment.title.split("|")[1])

      # Gets all the ids from the list and prepares them for fetching
      #print "Retrieving %d annotations" % (len(id_list))
      handle = Entrez.efetch("protein",id=",".join(self.gi_num), rettype="fasta", retmode="xml")

    except urllib2.HTTPError and ValueError as detail:
      self.success = False
      self.errors = detail
      return

    make_table_in = tempfile.NamedTemporaryFile()
    
    make_table_in.write("<TSeqSet>\n")
    # Parses the fetched records over its iterator
    for seq_record in Entrez.parse(handle):
        make_table_in.write("\t<TSeq>\n")
        make_table_in.write("\t\t<TSeq_seqtype value=\"protein\">protein</TSeq_seqtype>\n")
        if 'TSeq_gi' in seq_record:
            make_table_in.write("\t\t<TSeq_gi>" + seq_record['TSeq_gi'] + "</TSeq_gi>\n")
        if 'TSeq_accver' in seq_record:
            make_table_in.write("\t\t<TSeq_accver>" + seq_record['TSeq_accver'] + "</TSeq_accver>\n")
        if 'TSeq_taxid' in seq_record:
            make_table_in.write("\t\t<TSeq_taxid>" + seq_record['TSeq_taxid'] + "</TSeq_taxid>\n")
        if 'TSeq_orgname' in seq_record:
            make_table_in.write("\t\t<TSeq_orgname>" + seq_record['TSeq_orgname'] + "</TSeq_orgname>\n")
        if 'TSeq_defline' in seq_record:
            make_table_in.write("\t\t<TSeq_defline>" + seq_record['TSeq_defline'] + "</TSeq_defline>\n")
        if 'TSeq_length' in seq_record:
            make_table_in.write("\t\t<TSeq_length>" + seq_record['TSeq_length'] + "</TSeq_length>\n")
        if 'TSeq_sequence' in seq_record:
            make_table_in.write("\t\t<TSeq_sequence>" + seq_record['TSeq_sequence'] + "</TSeq_sequence>\n")
        make_table_in.write("\t</TSeq>\n")
    make_table_in.write("</TSeqSet>\n")

    # System call to make table
    process = subprocess.Popen("make_table5 -i " + make_table_in.name + " -c " + str(make_table_threshold), \
                               stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
    stdout, stderr = process.communicate()
    self.make_table_out = (stdout, stderr)
    # Get a file pointer to the faa file from make_table5
    self.faa_file = open(make_table_in.name + ".faa")
    self.success = True
    self.filepath = make_table_in.name + ".faa"
